import readlineSync from "readline-sync";
import { ChatOllama } from "@langchain/ollama";
import {
  ChatPromptTemplate,
  MessagesPlaceholder,
  SystemMessagePromptTemplate,
  HumanMessagePromptTemplate,
} from "@langchain/core/prompts";
import { RunnableWithMessageHistory } from "@langchain/core/runnables";
import { FileSystemChatMessageHistory } from "@langchain/community/stores/message/file_system";

// 模型
const model = new ChatOllama({
  model: "llama3",
  temperature: 0.7,
});

// 提示词
const prompt = ChatPromptTemplate.fromMessages([
  SystemMessagePromptTemplate.fromTemplate(`
    你是一名中文对话助手。无论用户使用何种语言，**必须始终用【简体中文】回答**。

    严格遵守：
    - 只用简体中文回答；除代码块外，禁止出现英文句子或段落。
    - 直接给出答案与步骤，避免客套和重复提示。
    - 列点请用精炼中文要点；必要时给出简短结论→理由→建议的结构。
    - 若需给代码，**代码内注释与说明也用中文**（保留必要的英文关键字/标识符即可）。
    - 问题含糊时，依据摘要与本轮上下文做最合理假设并继续回答，**不要反问**。
    
    示例（用于风格约束）：
    - ✅ “可以。建议你这样做：… 然后… 最后…”
    - ✅ “原因是：…；因此建议：…”
    - ❌ “Sure, here is …”
    - ❌ “Of course! …”
    `),
  new MessagesPlaceholder("history"), // 历史记录占位符
  HumanMessagePromptTemplate.fromTemplate("{input}"),
]);

// 配置对象
const cfg = {
  configurable: { sessionId: "demo-session" },
};

// 封装对话链
const chain = new RunnableWithMessageHistory({
  runnable: prompt.pipe(model), // 基础链
  getMessageHistory: async (sessionId) =>
    new FileSystemChatMessageHistory({
      sessionId,
      storageDir: "./history",
    }),
  inputMessagesKey: "input",
  historyMessagesKey: "history",
});

async function chatLoop() {
  console.log("开始对话，输入内容后回车；输入 /clear 清空历史，/exit 退出。");

  while (true) {
    // —— 用户输入 ——
    const input = readlineSync.question("\n你：").trim();
    if (!input) continue;

    // —— 命令处理 ——
    if (input === "/exit" || input === "exit" || input === "quit") {
      console.log("已退出。");
      break;
    }
    if (input === "/clear") {
      const h = await chain.getMessageHistory(cfg.configurable.sessionId);
      await h.clear();
      console.log("已清空历史。\n");
      continue;
    }

    try {
      const stream = await chain.stream({ input }, cfg);
      process.stdout.write("助理：");
      for await (const chunk of stream) {
        process.stdout.write(chunk.content);
      }
      console.log("\n");
    } catch (err) {
      console.error("调用模型失败：", err?.message ?? err);
      continue;
    }
  }
}

chatLoop();
